# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

import os
import sys
from pathlib import Path

import torch

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from det.models.common import DetectMultiBackend
from det.utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from det.utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from det.utils.torch_utils import select_device, time_sync


class Detector:
    @torch.no_grad()
    def __init__(self):
        # Load model
        weights = ROOT / 'best.pt'  # model.pt path(s)
        print(ROOT)
        data = ROOT / 'maskhelper.yaml'  # dataset.yaml path
        imgsz = (640, 640)  # inference size (height, width)
        self.device = select_device('')
        self.model = DetectMultiBackend(weights, device=self.device, dnn=False, data=data, fp16=False)
        self.stride, self.names, self.pt = self.model.stride, self.model.names, self.model.pt
        self.imgsz = check_img_size(imgsz, s=self.stride)  # check image size


    def detect(self,
               source=ROOT / 'data/images',  # file/dir/URL/glob, 0 for webcam,
               conf_thres=0.85,  # confidence threshold
               iou_thres=0.85,  # NMS IOU threshold
               max_det=10,  # maximum detections per image
               classes=None,  # filter by class: --class 0, or --class 0 2 3
               agnostic_nms=False,  # class-agnostic NMS
               augment=False,  # augmented inference
            ):
        self.source = str(source)
        print("1")
        # Dataloader
        dataset = LoadImages(self.source, img_size=self.imgsz, stride=self.stride, auto=self.pt)
        print(dataset)
        self.bs = 1  # batch_size
        # Run inference
        self.model.warmup(imgsz=(1 if self.pt else self.bs, 3, *self.imgsz))  # warmup
        seen, windows, dt = 0, [], [0.0, 0.0, 0.0]
        for path, im, im0s, vid_cap, s in dataset:
            t1 = time_sync()
            im = torch.from_numpy(im).to(self.device)
            im = im.half() if self.model.fp16 else im.float()  # uint8 to fp16/32
            im /= 255  # 0 - 255 to 0.0 - 1.0
            if len(im.shape) == 3:
                im = im[None]  # expand for batch dim
            t2 = time_sync()
            dt[0] += t2 - t1

            # Inference
            pred = self.model(im, augment=augment, visualize=False)
            t3 = time_sync()
            dt[1] += t3 - t2

            # NMS
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
            dt[2] += time_sync() - t3

            # Second-stage classifier (optional)
            # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

            # Process predictions
            for i, det in enumerate(pred):  # per image
                res = []
                seen += 1
                p, im0, frame = path, im0s.copy(), getattr(dataset, 'image', 0)

                s += '%gx%g ' % im.shape[2:]  # print string

                if len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                    # Write results
                    for *xyxy, conf, cls in reversed(det):
                        c = int(cls)  # integer class
                        label = self.names[c]
                        coords = [xyxy[0], xyxy[1], xyxy[2], xyxy[3]]

                        res.append([label,coords])
                print(res)
                return res

